package com.example.demo.ai.train;

import com.example.demo.ai.sftp.SftpService;
import com.jcraft.jsch.*;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.io.InputStream;

@Service
public class TrainService2 {

    @Autowired
    private SftpService sftpService;

    public String executeRemoteCommand(String yamlPath, String pyPath) throws JSchException {

        Channel channel = sftpService.createChannel("exec");
        ChannelExec exec = (ChannelExec) channel;
        if(exec == null) return null;

        String yoloDir = "/root/yolo-project";
        String activate = "source /root/yolo11/bin/activate";
        String command = String.format(
                "cd %s && %s && yolo detect train data= %s model= %s epochs=100 imgsz=640",
                yoloDir, activate,yamlPath,pyPath
        );

        exec.setCommand(command);

        try {
            // 获取输入流（命令输出）
            exec.setInputStream(null);
            exec.setErrStream(System.err);

            InputStream in = exec.getInputStream();
            exec.connect();

            // 读取命令输出
            byte[] tmp = new byte[1024];
            while (true) {
                while (in.available() > 0) {
                    int i = in.read(tmp, 0, 1024);
                    if (i < 0) break;
                    System.out.print(new String(tmp, 0, i));
                }
                if (channel.isClosed()) {
                    if (in.available() > 0) continue;
                    System.out.println("Exit status: " + channel.getExitStatus());
                    break;
                }
                Thread.sleep(1000);
            }
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            // 关闭连接
            exec.disconnect();
        }

        return null;
    }
}
